%% Read audio file
% Note: due to internal regulations, we are not releasing audio recording data.
% Instead, we are providing the resulting time correspondence from temporal
% alignment of the audio files + the original code reading the
% audio files + screenshots of the audio files in waveforms.

% For subject C001, the speech sample number is 3
% For subject C002, the speech sample number is 4
% For subject C003, the speech sample number is 3
% For subject C004, the speech sample number is 3
sub = 'C001'; sample = '3';
[aud,Fs] = audioread(['/sampleAtlas/' sub '/sample' sample '.wav']); % we are not providing this file
audioRate = Fs;
imRate = 40;
figure; plot(aud(:,1));

%% Pick a stimulus repetition from the audio
int1 = 177000; int2 = 188000;
figure; plot(aud(int1:int2,1)); set(gcf,'color','w'); set(gca,'YTickLabel',[]);
ylabel(sub)
sound(aud(int1:int2,1),audioRate);
oneRep = aud(int1:int2,1);
oneRepL = length(oneRep);

%% Determine the image time frames from the stimulus intervals
peakIndEnd = int2;
peakIndBeg = peakIndEnd - oneRepL + 1;
peakInd = [peakIndBeg, peakIndEnd];
timesep = round(peakInd/audioRate*imRate);
disp(timesep)

%% Repeat the same process for all subjects, below are the choices used in this example package
% C001 image time frames are 443-470
% C002 image time frames are 975-1003
% C003 image time frames are 1053-1080
% C004 image time frames are 1080-1108

%% Create the template subject using Sub 1, rename the time frames
% Note: this example code is using the Matlab Tools for NIfTI and ANALYZE image Version 1.27.0.0 (426 KB) by Jimmy Shen
% Please install it from https://www.mathworks.com/matlabcentral/fileexchange/8797-tools-for-nifti-and-analyze-image
% load_untouch_nii and save_untouch_nii can be replaced by Matlab's niftiread and niftiwrite
c = 1;
for t = 443:470
    nii = load_untouch_nii(['/sampleAtlas/C001/sample' sample '/DynSp_MGH_C001_20220415_sample' sample '_Frame' num2str(t,'%05.f') '.nii']);
    save_untouch_nii(nii,['/sampleAtlas/C001/sample' sample '/C001_sample' sample '_aligned_tf' num2str(c,'%02.f') '.nii']);
    c = c+1;
end

%% Rename all other subjects' time frame number from 1 to T
% For C002:
c = 1;
for t = 975:1003
    nii = load_untouch_nii(['/sampleAtlas/C002/sample' sample '/DynSp_MGH_C002_20220624_sample' sample '_Frame' num2str(t,'%05.f') '.nii']);
    save_untouch_nii(nii,['/sampleAtlas/C002/sample' sample '/C002_sample' sample '_beforeAligned_tf' num2str(c,'%02.f') '.nii']);
    c = c+1;
end

%% Rename all other subjects' time frame number from 1 to T
% For C003:
c = 1;
for t = 1053:1080
    nii = load_untouch_nii(['/sampleAtlas/C003/sample' sample '/DynSp_MGH_C003_20220824_sample' sample '_Frame' num2str(t,'%05.f') '.nii']);
    save_untouch_nii(nii,['/sampleAtlas/C003/sample' sample '/C003_sample' sample '_beforeAligned_tf' num2str(c,'%02.f') '.nii']);
    c = c+1;
end

%% Rename all other subjects' time frame number from 1 to T
% For C004:
c = 1;
for t = 1080:1108
    nii = load_untouch_nii(['/sampleAtlas/C004/sample' sample '/DynSp_MGH_C004_20221212_sample' sample '_Frame' num2str(t,'%05.f') '.nii']);
    save_untouch_nii(nii,['/sampleAtlas/C004/sample' sample '/C004_sample' sample '_beforeAligned_tf' num2str(c,'%02.f') '.nii']);
    c = c+1;
end

%% Create a time table matching the repetition's sound waveforms
% Each row is a subject, the first subject is the reference template
timeMatch = [
    1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28;
    1.4 2.5 3.6 4.7 5.8 6.9 7.5 8.1 9.1 10.2 11.4 12.5 13.6 14.7 15.4 16.9 17.4 20.3 21.1 21.9 22.8 23.7 24.6 25.5 26.4 27.3 28.2 28.9;
    2.2 3.0 3.8 4.6 5.4 6.2 7.0 8.8 9.6 10.4 11.0 11.7 12.4 13.2 14.1 15.0 15.4 17.6 18.5 19.4 20.3 21.2 22.1 22.9 23.8 24.7 25.6 26.4;
    1.8 2.5 3.6 4.7 5.8 6.8 7.8 8.8 9.8 10.7 11.8 12.8 13.8 14.8 15.7 16.9 18.0 19.2 20.0 21.1 22.1 23.0 24.1 25.2 26.0 27.0 28.2 29];

%% Generate a prompt for ANTs to perform registration on and find the motion field between neighboring frames
% Take subject C002 for example:
for i = 1:size(timeMatch,2)
    t = timeMatch(2,i);
    if round(t) == t
        nii = load_untouch_nii(['/sampleAtlas/C002/sample' sample '/C002_sample' sample '_beforeAligned_tf' num2str(t,'%02.f') '.nii']);
        save_untouch_nii(nii,['/sampleAtlas/C002/sample' sample '/C002_sample' sample '_aligned_tf' num2str(i,'%02.f') '.nii']);
    else
        t1 = floor(t);
        t2 = ceil(t);
        antsPrompt = ["ANTS 3 -m 'CC[/sampleAtlas/C002/sample" sample "/C002_sample" sample "_beforeAligned_tf" num2str(t2,'%02.f') ".nii," + ...
            "/sampleAtlas/C002/sample" sample "/C002_sample" sample "_beforeAligned_tf" num2str(t1,'%02.f') ".nii,1,4]' " + ...
            "-o '/sampleAtlas/C002/sample" sample "/field" num2str(i,'%02.f') "_' -i 100x80x50 -r 'Gauss[3,0]' -t 'SyN[0.25]'"];
        antsPrompt = join(antsPrompt,'');
        disp(antsPrompt)
    end
end

%% After running the prompts in ANTs, motion fields will be generated
% Trim the motion fields
% Take subject C002 for example:
for i = 1:size(timeMatch,2)
    t = timeMatch(2,i);
    if round(t) == t
        continue;
    else
        proportionMotion = t - floor(t);
        nii = load_untouch_nii(['/sampleAtlas/C002/sample' sample '/field' num2str(i,'%02.f') '_Warp.nii']);
        nii.img = proportionMotion*nii.img;
        save_untouch_nii(nii,['/sampleAtlas/C002/sample' sample '/field' num2str(i,'%02.f') '_Warp.nii']);
    end
end

%% Finally, go back to ANTs and use the trimmed motion fields to warp the images
% Generate the ANTs prompt
% Take subject C002 for example:
for i = 1:size(timeMatch,2)
    t = timeMatch(2,i);
    if round(t) == t
        continue;
    else
        t1 = floor(t);
        t2 = ceil(t);
        antsPrompt = ["WarpImageMultiTransform 3 '/sampleAtlas/C002/sample" sample "/C002_sample" sample "_beforeAligned_tf" num2str(t1,'%02.f') ".nii' " + ...
            "'/sampleAtlas/C002/sample" sample "/C002_sample" sample "_aligned_tf" num2str(i,'%02.f') ".nii' " + ...
            "-R '/sampleAtlas/C002/sample" sample "/C002_sample" sample "_beforeAligned_tf" num2str(t2,'%02.f') ".nii' " + ...
            "'/sampleAtlas/C002/sample" sample "/field" num2str(i,'%02.f') "_Warp.nii' " + ...
            "'/sampleAtlas/C002/sample" sample "/field" num2str(i,'%02.f') "_Affine.nii'"];
        antsPrompt = join(antsPrompt,'');
        disp(antsPrompt)
    end
end

%% The final time aligned results are named '/sampleAtlas/C00I/sampleJ/C00I_sampleJ_aligned_tfT.nii'

%% From now on, we are going to combine each time frame into an atlas 3D volume for each time frame
% Create a separate folder for each of time frame 1 to 28, named "1",
% "2","3",... and put all subjects' corresponding frame inside


% Run the following parts in Python
% The "paste" executable is compiled for a MAC system, please use a MAC for
% this zero padding process


% Run run_paste.py for zero padding and create the _p_aligned_tf files


% Create an "atlas" folder within each time frame folder and copy and paste
% all _p_aligned_tf files inside


% Run the run_atlas.sh file to create the atlas, it should be named as
% xxxxx_template.nii.gz (please change output name as needed).


% After you've done the above on time frame 1, copy and paste all Warp and
% Affine files and tf1_xxxxx_template.nii.gz to the parent folder


% Finally run run_reg.py for transform of all remaining frames using the transformations learned during the 1st frame's atlas construction process


% Copy and paste the_registered files into the separate time frame folder
% and run run_atlas.sh for all time frames individually, then you will get the final atlas for each time frame. You can combine those indivdual time frames into a single 4D atlas using your favorite tools.



